Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements a comprehensive "Strict Resume" feature for Transformers models, enabling the restoration of full training state including optimizer, scheduler, scaler, RNG states, and data progress. Key changes involve implementing load_training_state and read_training_progress across the model, server, and client layers, alongside dataloader enhancements to support sample-level skipping for map-style datasets. Feedback highlights several critical improvements: ensuring deterministic RNG in distributed settings by avoiding unseeded random states, addressing the deprecated use of StopIteration in generators, improving security by using weights_only=True during checkpoint loading, and removing an accidental BOM character in the client generator. Additionally, a more robust approach for re-initializing the dataloader is suggested to avoid modifying private PyTorch attributes.
| self.dataloader.__initialized = False | ||
| self._rebuild_sampler_stack() | ||
| self.dataloader.__initialized = True |
There was a problem hiding this comment.
Accessing and modifying the private attribute __initialized of torch.utils.data.DataLoader is brittle and relies on internal implementation details of PyTorch that could change. A safer approach to update the sampler stack after the dataloader has been created is to simply re-instantiate the underlying self.dataloader using the stored self.dataloader_params.
| self.dataloader.__initialized = False | |
| self._rebuild_sampler_stack() | |
| self.dataloader.__initialized = True | |
| if self.dataloader is not None: | |
| self.dataloader = None | |
| self._lazy_init_dataloader() |
|
/gemini summary |
Summary of ChangesThis pull request introduces robust checkpoint resumption capabilities to the training framework. By enabling the restoration of full training states—including optimizer, scheduler, and RNG configurations—and implementing precise data skipping in the dataloader, the changes ensure that training can be reliably resumed after interruptions. Additionally, the PR optimizes checkpoint handling for FSDP2 strategies and adds necessary API endpoints to support these features in distributed and remote training environments. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Activity
|
1 participant
PR type
PR information
在TrasnfomersModel和MultiLoraModel实现完整训练状态的恢复——包括优化器、调度器、RNG配置以及数据集跳过